import socket as sk
import _thread
import time
from tools.transmit_receive.receiver import Receiver
from tools.o_gan_module import TrafficModelUser, TrafficModelTrainner, TrafficModelParam
import os

LISTEN_SIZE = 2
BASE = 1

EPOCHS = 10
BATCH_SIZE = 64
NOISE_DIM = 8
TRAFFIC_DIM = 15
G_PATH = 'model/g_model.h5'
D_PATH = 'model/d_model.h5'


def get_or_create_gan_model(g_path: str, d_path: str) -> TrafficModelUser:
    """
    创建一个模型使用者
    搭配上一个训练器
    :param g_path:
    :param d_path:
    :return:
    """
    param = TrafficModelParam(TRAFFIC_DIM, NOISE_DIM)
    # model_trainer = TrafficModelTrainner(EPOCHS, BATCH_SIZE, NOISE_DIM, TRAFFIC_DIM)
    return TrafficModelUser(g_path, d_path, param)


def main():
    host_socket = sk.socket(sk.AF_INET, sk.SOCK_STREAM)
    host_socket.bind(('', 6666))
    host_socket.listen(LISTEN_SIZE)
    model_user = get_or_create_gan_model(os.path.join(os.getcwd(), G_PATH),
                                         os.path.join(os.getcwd(), D_PATH))

    while True:
        client_socket, client_addr = host_socket.accept()
        _thread.start_new_thread(listen_thread, (client_socket, model_user,))


def listen_thread(client_socket: sk.socket, model_user: TrafficModelUser):
    receiver = Receiver(BASE)
    while True:
        data = client_socket.recv(4096)
        tick = time.time()
        if len(data) == 0:
            break
        print(data)
        print(tick)
        receiver.add_tick(tick)
    restored_noise = model_user.analysis(receiver.answer())
    client_socket.close()
    print(restored_noise)


if __name__ == '__main__':
    main()
